import os

from PyQt5.QtCore import QTimer
from PyQt5.QtWidgets import QApplication, QWidget, QGridLayout, QLabel, QPushButton, QHBoxLayout, QSizePolicy, QLineEdit
from PyQt5.QtGui import QPixmap, QFont
from PyQt5 import QtCore

class SameNameImageViewer(QWidget):
    size=400
    def __init__(self, parent_folder):
        super().__init__()
        self.parent_folder = parent_folder
        self.image_groups = {}
        self.current_group_index = 0
        self.initUI()

    def initUI(self):
        self.find_same_named_images()
        self.display_images()
        self.create_buttons()
        self.create_image_input()
        self.setMinimumSize(400, 300)
        self.setWindowTitle('同名图片查看器')
        self.show()

    def create_image_input(self):
        self.input_layout = QHBoxLayout()
        self.image_input = QLineEdit()
        self.input_layout.addWidget(self.image_input)
        self.show_button = QPushButton("显示图片")
        self.show_button.clicked.connect(self.show_image_by_name)
        self.input_layout.addWidget(self.show_button)
        self.layout.addLayout(self.input_layout, len(self.image_groups) * 2 + 2, 0, 1, 4)

    def show_image_by_name(self):
        image_name = self.image_input.text()
        print(image_name)
        if image_name in self.image_groups:
            if self.layout is not None:
                for i in reversed(range(self.layout.count())):
                    widget = self.layout.itemAt(i).widget()
                    if widget is not None:
                        widget.setParent(None)

            row, col = 0, 0
            group_name = image_name
            print("group_name", group_name)
            group_label = QLabel(f"图: {group_name}")
            group_label.setAlignment(QtCore.Qt.AlignCenter)  # 设置组名居中对齐
            self.layout.addWidget(group_label, row, col, 1, 4)
            row += 1
            col = 0
            for image_info in self.image_groups[group_name]:
                image_path, parent_folder = image_info
                image_label = QLabel()
                # image_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)  # 设置图片标签的大小策略
                image_label.setFixedSize(self.size, self.size)
                pixmap = QPixmap(image_path)
                pixmap = pixmap.scaled(self.size, self.size, QtCore.Qt.KeepAspectRatio, QtCore.Qt.SmoothTransformation)
                image_label.setPixmap(pixmap)
                # image_label.setScaledContents(True)  # 将 pixmap 缩放以适应 QLabel
                self.layout.addWidget(image_label, row, col)

                font = QFont()
                font.setPointSize(8)
                folder_label = QLabel(parent_folder)
                folder_label.setAlignment(QtCore.Qt.AlignCenter)  # 设置组名居中对齐
                folder_label.setFont(font)  # 设置字体大小

                self.layout.addWidget(folder_label, row + 1, col)
                col = (col + 1) % 4
                if col == 0:
                    row += 2

            self.setLayout(self.layout)

    def find_same_named_images(self):

        # list
        folder_names = [
            'image',
            'label',
            'U_Net',
            'SegNet',
            'NestedUNet',
            'PSPNet',
            # 'LiteST_Net',
            # 'LGFO_Net'
            # 'A',
            # 'B',
            # 'C',
            # 'D',
            # 'predict'，
            'VIT'
        ]

        # 遍历列表中的文件夹名称，依次处理
        for folder_name in folder_names:
            folder_path = os.path.join(self.parent_folder, folder_name, 'predictResult')
            # 检查文件夹是否存在
            if not os.path.exists(folder_path):
                continue
            # 调用处理文件夹的方法
            self.process_folder(folder_path)


    def process_folder(self, folder_path):
        for file in os.listdir(folder_path):
            if file.endswith('.png') or file.endswith('.jpg'):
                image_path = os.path.join(folder_path, file)
                image_name = os.path.splitext(file)[0]
                parent_folder = os.path.basename(os.path.dirname(folder_path))
                if image_name not in self.image_groups:
                    self.image_groups[image_name] = []
                self.image_groups[image_name].append((image_path, parent_folder))

    def display_images(self):
        self.layout = QGridLayout(self)
        self.layout.setSpacing(5)  # 设置布局中组件之间的间距为 10 像素
        self.update_display()

    def update_display(self):
        if self.layout is not None:
            for i in reversed(range(self.layout.count())):
                widget = self.layout.itemAt(i).widget()
                if widget is not None:
                    widget.setParent(None)

        row, col = 0, 0
        group_name = list(self.image_groups.keys())[self.current_group_index]
        group_label = QLabel(f"图: {group_name}")
        group_label.setAlignment(QtCore.Qt.AlignCenter)  # 设置组名居中对齐
        self.layout.addWidget(group_label, row, col, 1, 4)
        row += 1
        col = 0
        for image_info in self.image_groups[group_name]:
            image_path, parent_folder = image_info
            image_label = QLabel()
            # image_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)  # 设置图片标签的大小策略
            image_label.setFixedSize(self.size,self.size)
            pixmap = QPixmap(image_path)
            pixmap = pixmap.scaled(self.size, self.size, QtCore.Qt.KeepAspectRatio, QtCore.Qt.SmoothTransformation)
            image_label.setPixmap(pixmap)
            # image_label.setScaledContents(True)  # 将 pixmap 缩放以适应 QLabel
            self.layout.addWidget(image_label, row, col)

            font=QFont()
            font.setPointSize(8)
            folder_label = QLabel(parent_folder)
            folder_label.setAlignment(QtCore.Qt.AlignCenter)  # 设置组名居中对齐
            folder_label.setFont(font)   # 设置字体大小

            self.layout.addWidget(folder_label, row + 1, col)
            col = (col + 1) % 4
            if col == 0:
                row += 2

        self.setLayout(self.layout)

    def create_buttons(self):
        button_layout = QHBoxLayout()

        prev_button = QPushButton("上一组")
        prev_button.clicked.connect(self.previous_group)
        button_layout.addWidget(prev_button)

        next_button = QPushButton("下一组")
        next_button.clicked.connect(self.next_group)
        button_layout.addWidget(next_button)

        # next_button = QPushButton("调整图片大小")
        # next_button.clicked.connect(self.change)
        # button_layout.addWidget(next_button)

        self.layout.addLayout(button_layout, len(self.image_groups) * 2 + 1, 0, 1, 4)

    def previous_group(self):
        if self.current_group_index > 0:
            self.current_group_index -= 1
            self.update_display()

    def next_group(self):
        if self.current_group_index < len(self.image_groups) - 1:
            self.current_group_index += 1
            self.update_display()

    def change(self):
        self.update_display()




if __name__ == '__main__':
    app = QApplication([])
    parent_folder = '/media/jolly/HIK/CorsicanFire/BoFire'  # 替换为您的父文件夹路径
    window = SameNameImageViewer(parent_folder)
    app.exec_()